import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

# 创建复平面网格
x = np.linspace(-2, 2, 500)
y = np.linspace(-2, 2, 500)
X, Y = np.meshgrid(x, y)
Z = X + 1j * Y

# 计算主分支平方根
with np.errstate(invalid='ignore'):
    W_principal = np.sqrt(Z)

# 提取实部和虚部
real_principal = np.real(W_principal)
imag_principal = np.imag(W_principal)

# 创建3D图形
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 归一化颜色映射
norm = plt.Normalize(imag_principal.min(), imag_principal.max())

# 绘制主分支曲面
surf = ax.plot_surface(
    X, Y, real_principal, 
    facecolors=cm.jet(norm(imag_principal)),
    rstride=2, cstride=2,
    alpha=0.9, antialiased=True
)

# 设置标签和标题
ax.set_xlabel('Re(z)', fontsize=12)
ax.set_ylabel('Im(z)', fontsize=12)
ax.set_zlabel('Re($\sqrt{z}$)', fontsize=12)
ax.set_title('Principal Branch of $w = \sqrt{z}$', fontsize=16)

# 添加颜色条
mappable = cm.ScalarMappable(cmap=cm.jet, norm=norm)
mappable.set_array(imag_principal)
cbar = fig.colorbar(mappable, ax=ax, shrink=0.7)
cbar.set_label('Im($\sqrt{z}$)', fontsize=12)

# 设置视角
ax.view_init(elev=30, azim=-60)

plt.tight_layout()
plt.savefig('principal_branch.png', dpi=300)
plt.show()

# 创建更精细的网格以展示多值性
theta = np.linspace(0, 4*np.pi, 500)  # 两圈以展示多值性
r = np.linspace(0, 2, 500)
R, Theta = np.meshgrid(r, theta)

# 转换为笛卡尔坐标
X_riemann = R * np.cos(Theta)
Y_riemann = R * np.sin(Theta)
Z_riemann = X_riemann + 1j * Y_riemann

# 计算平方根（考虑多值性）
W_riemann = np.sqrt(R) * np.exp(1j * Theta/2)

# 提取实部和虚部
real_riemann = np.real(W_riemann)
imag_riemann = np.imag(W_riemann)

# 创建3D图形
fig = plt.figure(figsize=(14, 12))
ax = fig.add_subplot(111, projection='3d')

# 归一化颜色映射
norm_riemann = plt.Normalize(imag_riemann.min(), imag_riemann.max())

# 绘制黎曼面
surf = ax.plot_surface(
    X_riemann, Y_riemann, real_riemann, 
    facecolors=cm.jet(norm_riemann(imag_riemann)),
    rstride=4, cstride=4,
    alpha=0.85, antialiased=True
)

# 设置标签和标题
ax.set_xlabel('Re(z)', fontsize=12)
ax.set_ylabel('Im(z)', fontsize=12)
ax.set_zlabel('Re($\sqrt{z}$)', fontsize=12)
ax.set_title('Riemann Surface of $w = \sqrt{z}$ (Showing Multi-valuedness)', fontsize=16)

# 添加颜色条
mappable_riemann = cm.ScalarMappable(cmap=cm.jet, norm=norm_riemann)
mappable_riemann.set_array(imag_riemann)
cbar = fig.colorbar(mappable_riemann, ax=ax, shrink=0.7)
cbar.set_label('Im($\sqrt{z}$)', fontsize=12)

# 设置视角
ax.view_init(elev=25, azim=-50)

plt.tight_layout()
plt.savefig('riemann_surface.png', dpi=300)
plt.show()